/**
  * Copyright (c) CSIRO Australia, 2009
  *
  * @author $Author: jia020 $
  * @version $Id: XmlSignatureUtil.java 643 2009-05-20 00:20:55Z jia020 $
  */
package au.csiro.nt.pdsp.util;

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.StringReader;
import java.security.InvalidAlgorithmParameterException;
import java.security.Key;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.UnrecoverableEntryException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;

import javax.xml.crypto.AlgorithmMethod;
import javax.xml.crypto.KeySelector;
import javax.xml.crypto.KeySelectorException;
import javax.xml.crypto.KeySelectorResult;
import javax.xml.crypto.MarshalException;
import javax.xml.crypto.XMLCryptoContext;
import javax.xml.crypto.XMLStructure;
import javax.xml.crypto.dsig.CanonicalizationMethod;
import javax.xml.crypto.dsig.DigestMethod;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.Transform;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureException;
import javax.xml.crypto.dsig.XMLSignatureFactory;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.dom.DOMValidateContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;
import javax.xml.parsers.ParserConfigurationException;
import org.w3c.dom.Document;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;

public class XmlSignatureUtil
{

	public static String signFile(String theKeyStorePath,
			String theKeyStorePasswd, String theKeyName,
			String theInputFileName, String theOutputFileName)
	{
		try
		{
			Document doc = XmlUtil.fileToDocument(theInputFileName);
			return signDocument(theKeyStorePath, theKeyStorePasswd, theKeyName,
					doc);

		} catch (FileNotFoundException e)
		{
			e.printStackTrace();
		} catch (IOException e)
		{
			e.printStackTrace();
		} catch (SAXException e)
		{
			e.printStackTrace();
		} catch (ParserConfigurationException e)
		{
			e.printStackTrace();
		}
		return null;
	}

	public static String signString(String theKeyStorePath,
			String theKeyStorePasswd, String theKeyName, String theContent)
	{
		try
		{
			Document doc = XmlUtil.stringToDocument(theContent);
			return signDocument(theKeyStorePath, theKeyStorePasswd, theKeyName,
					doc);
		} catch (FileNotFoundException e)
		{
			e.printStackTrace();
		} catch (IOException e)
		{
			e.printStackTrace();
		} catch (SAXException e)
		{
			e.printStackTrace();
		} catch (ParserConfigurationException e)
		{
			e.printStackTrace();
		}
		return null;
	}

	public static Document signDoc(String theKeyStorePath,
			String theKeyStorePasswd, String theKeyName, Document theDoc)
	{
		// Create a DOM XMLSignatureFactory that will be used to
		// generate the enveloped signature.
		XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");

		// Create a Reference to the enveloped document (in this case,
		// you are signing the whole document, so a URI of "" signifies
		// that, and also specify the SHA1 digest algorithm and
		// the ENVELOPED Transform.
		Reference ref;
		try
		{
			ref = fac.newReference("", fac.newDigestMethod(DigestMethod.SHA1,
					null), Collections.singletonList(fac.newTransform(
					Transform.ENVELOPED, (TransformParameterSpec) null)), null,
					null);
			SignedInfo si = fac.newSignedInfo(fac.newCanonicalizationMethod(
					CanonicalizationMethod.INCLUSIVE,
					(C14NMethodParameterSpec) null), fac.newSignatureMethod(
					SignatureMethod.RSA_SHA1, null), Collections
					.singletonList(ref));

			// Load the KeyStore and get the signing key and certificate.
			KeyStore ks = KeyStore.getInstance("JKS");
			ks.load(new FileInputStream(theKeyStorePath), theKeyStorePasswd
					.toCharArray());
			KeyStore.PrivateKeyEntry keyEntry = (KeyStore.PrivateKeyEntry) ks
					.getEntry(theKeyName, new KeyStore.PasswordProtection(
							theKeyStorePasswd.toCharArray()));
			X509Certificate cert = (X509Certificate) keyEntry.getCertificate();

			// Create the KeyInfo containing the X509Data.
			KeyInfoFactory kif = fac.getKeyInfoFactory();
			List x509Content = new ArrayList();
			x509Content.add(cert.getSubjectX500Principal().getName());
			x509Content.add(cert);
			X509Data xd = kif.newX509Data(x509Content);
			KeyInfo ki = kif.newKeyInfo(Collections.singletonList(xd));

			// Create a DOMSignContext and specify the RSA PrivateKey and
			// location of the resulting XMLSignature's parent element.
			DOMSignContext dsc = new DOMSignContext(keyEntry.getPrivateKey(),
					theDoc.getDocumentElement());

			// Create the XMLSignature, but don't sign it yet.
			XMLSignature signature = fac.newXMLSignature(si, ki);

			// Marshal, generate, and sign the enveloped signature.
			signature.sign(dsc);

			return theDoc;
		} catch (NoSuchAlgorithmException e)
		{
			e.printStackTrace();
		} catch (InvalidAlgorithmParameterException e)
		{
			e.printStackTrace();
		} catch (KeyStoreException e)
		{
			e.printStackTrace();
		} catch (CertificateException e)
		{
			e.printStackTrace();
		} catch (FileNotFoundException e)
		{
			e.printStackTrace();
		} catch (IOException e)
		{
			e.printStackTrace();
		} catch (UnrecoverableEntryException e)
		{
			e.printStackTrace();
		} catch (MarshalException e)
		{
			e.printStackTrace();
		} catch (XMLSignatureException e)
		{
			e.printStackTrace();
		}
		return null;
	}

	public static String signDocument(String theKeyStorePath,
			String theKeyStorePasswd, String theKeyName, Document theDoc)
	{
		Document signed_doc = signDoc(theKeyStorePath, theKeyStorePasswd,
				theKeyName, theDoc);

		if (signed_doc == null)
			return null;

		return XmlUtil.documentToString(signed_doc);
	}

	public static int validateDocument(Document theDoc) throws Exception
	{
		XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");
		// Find Signature element.
		NodeList nl = theDoc.getElementsByTagNameNS(XMLSignature.XMLNS,
				"Signature");
		if (nl.getLength() == 0)
		{
			throw new Exception("Cannot find Signature element");
		}

		// Create a DOMValidateContext and specify a KeySelector
		// and document context.
		DOMValidateContext valContext = new DOMValidateContext(
				new X509KeySelector(), nl.item(0));

		// Unmarshal the XMLSignature.
		XMLSignature signature = fac.unmarshalXMLSignature(valContext);

		// Validate the XMLSignature.
		boolean coreValidity = signature.validate(valContext);
		// Check core validation status.
		if (coreValidity == false)
		{
			System.err.println("Signature failed core validation");
			boolean sv = signature.getSignatureValue().validate(valContext);
			System.out.println("signature validation status: " + sv);
			if (sv == false)
			{
				// Check the validation status of each Reference.
				Iterator i = signature.getSignedInfo().getReferences()
						.iterator();
				for (int j = 0; i.hasNext(); j++)
				{
					boolean refValid = ((Reference) i.next())
							.validate(valContext);
					System.out.println("ref[" + j + "] validity status: "
							+ refValid);
				}
			}
			return -1;
		} else
		{
			System.out.println("Signature passed core validation");
		}
		return 0;
	}
}
